
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.neighbors import KNeighborsClassifier
from sklearn.svm import SVC
from sklearn.ensemble import RandomForestClassifier  # 新增随机森林导入
from sklearn.metrics import accuracy_score, classification_report
from joblib import dump

class ClassificationModel(object):
    def __init__(self):
        self.df = pd.read_csv('fina_indicator.csv')

    def get_conditions(self):
        # 分类前准备
        df = self.df
        df['max_ratio'] = df['max_closes'] / df['the_closes']  # 收益潜力
        df['min_ratio'] = df['min_closes'] / df['the_closes']  # 风险程度

        # 划分高低阈值
        high_retuen_threshold = df['max_ratio'].quantile(0.4)  # 前40% 定义为高收益
        high_risk_threshold = df['min_ratio'].quantile(0.4)  # 前40% 定义为高风险
        # 生成分类标签
        conditions = [
            (df['max_ratio'] >= high_retuen_threshold) & (df['min_ratio'] >= high_risk_threshold),
            (df['max_ratio'] < high_retuen_threshold) & (df['min_ratio'] >= high_risk_threshold),
            (df['max_ratio'] >= high_retuen_threshold) & (df['min_ratio'] < high_risk_threshold),
            (df['max_ratio'] < high_retuen_threshold) & (df['min_ratio'] < high_risk_threshold)
        ]
        labels = ['高收益高风险', '低收益高风险', '高收益低风险', '低收益低风险']
        df['category'] = np.select(conditions, labels, default='未知')
        # 标签编码
        label_encoder = LabelEncoder()
        df['category_encoded'] = label_encoder.fit_transform(df['category'])
        # 数据标准化
        scaler = StandardScaler()
        features = df[['eps', 'total_revenue_ps', 'undist_profit_ps', 'gross_margin', 'fcff', 'fcfe',
                       'tangible_asset', 'bps', 'grossprofit_margin', 'npta']]
        X = scaler.fit_transform(features)
        y = df['category_encoded']
        # 划分训练集和测试集
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=24)
        dump(scaler, 'scaler.jobilb')
        dump(label_encoder,'label_encoder.jobilb')
        return X_train, X_test, y_train, y_test, label_encoder

    def knn_utils(self, X_train, X_test, y_train, y_test, label_encoder):
        '''
        KNN 分类
        '''
        # 初始化 KNN 分类器
        knn = KNeighborsClassifier(n_neighbors=5)
        # 训练模型
        knn.fit(X_train, y_train)
        # 预测测试集
        y_pred = knn.predict(X_test)
        # 评估模型
        accuracy = accuracy_score(y_test, y_pred)
        report = classification_report(y_test, y_pred, target_names=label_encoder.classes_)
        print(f"KNN 模型准确率: {accuracy:.4f}")
        print("KNN 分类报告:")
        print(report)
        dump(knn, 'knn.jobilb')

    def svc_utils(self, X_train, X_test, y_train, y_test, label_encoder):
        '''
        SVC 分类
        '''
        # 初始化 SVC 分类器
        svc = SVC(kernel='linear', random_state=24)
        # 训练模型
        svc.fit(X_train, y_train)
        # 预测测试集
        y_pred = svc.predict(X_test)
        # 评估模型
        accuracy = accuracy_score(y_test, y_pred)
        report = classification_report(y_test, y_pred, target_names=label_encoder.classes_)
        print(f"SVC 模型准确率: {accuracy:.4f}")
        print("SVC 分类报告:")
        print(report)
        dump(svc, 'svc.jobilb')

    # 新增随机森林分类器方法
    def rf_utils(self, X_train, X_test, y_train, y_test, label_encoder):
        '''
        随机森林分类
        '''
        # 初始化随机森林分类器
        rf = RandomForestClassifier(n_estimators=100, random_state=24)  # 100棵树，固定随机种子
        # 训练模型
        rf.fit(X_train, y_train)
        # 预测测试集
        y_pred = rf.predict(X_test)
        # 评估模型
        accuracy = accuracy_score(y_test, y_pred)
        report = classification_report(y_test, y_pred, target_names=label_encoder.classes_)
        print(f"随机森林模型准确率: {accuracy:.4f}")
        print("随机森林分类报告:")
        print(report)
        dump(rf, 'rf.jobilb')  # 保存模型

if __name__ == '__main__':
    cu = ClassificationModel()
    X_train, X_test, y_train, y_test, label_encoder = cu.get_conditions()
    cu.knn_utils(X_train, X_test, y_train, y_test, label_encoder)
    cu.svc_utils(X_train, X_test, y_train, y_test, label_encoder)
    cu.rf_utils(X_train, X_test, y_train, y_test, label_encoder)  # 调用随机森林分类器